from PIL import Image, ImageSequence
import threading
import queue
import glob
import os
import requests
from hashlib import md5
from PIL import ImageDraw, ImageFont


class GifHandler:

    def __init__(self, output_folder,watermark=None):
        self.output_folder = output_folder
        self.frame_queue = queue.PriorityQueue()
        self.watermark = watermark

    def url_to_md5(self, url):
        return md5(url.encode('utf-8')).hexdigest()

    def download_file(self, url, local_path):
        with requests.get(url, stream=True) as r:
            r.raise_for_status()
            with open(local_path, 'wb') as f:
                for chunk in r.iter_content(chunk_size=8192): 
                    f.write(chunk)
    
    def download_and_split_gif(self, url, force_update=False):
        filename = self.url_to_md5(url)
        file_dir = os.path.join(self.output_folder, filename)

        if not os.path.exists(file_dir):
            os.makedirs(file_dir)

        gif_local_path = os.path.join(file_dir, filename + '.gif')

        if not os.path.isfile(gif_local_path) or force_update:
            self.download_file(url, gif_local_path)
        
        # 这里有分割了!那就配置成面向对象的方式
        frame_delays = self.split_gif(gif_local_path)
        frames = [os.path.join(file_dir, f"frame_{i}.png") for i in range(self.frame_count(gif_local_path))]
        os.remove(gif_local_path)
        # 这里需要多回传一个gif的原始数据
        return frames,file_dir , frame_delays

    def frame_count(self, gif_path):
        with Image.open(gif_path) as img:
            return sum(1 for _ in ImageSequence.Iterator(img))

    def split_gif(self, gif_path):
        with Image.open(gif_path) as im:
            threads = []
            delays = []
            for i, frame in enumerate(ImageSequence.Iterator(im)):
                delays.append(frame.info['duration'])
                frame_dir = os.path.dirname(gif_path)
                frame_path = os.path.join(frame_dir, f"frame_{i}.png")  # 保存成png的模式(这里是用多线程去下载的)
                self.frame_queue.put((i, frame_path))
                thread = threading.Thread(target=self.save_frame, args=(frame.copy(), i, frame_dir))
                thread.start()
                threads.append(thread)
            for thread in threads:
                thread.join()
            return delays

    def save_frame(self, frame, frame_number, frame_dir):
        watermark = self.watermark
        if watermark is not None:
            width, height = frame.size
            font_size = int(width * watermark.get('font_scale', 1/20))  # Set font size to be a fraction of the image width
            script_dir = os.path.dirname(os.path.realpath(__file__))
            fontfile_path = os.path.join(script_dir, '..', 'resources', 'LXGWHeartSerifCHS.ttf')
            font = ImageFont.truetype(fontfile_path, font_size)
            offset = int(height * watermark.get('offset_ratio', 0))  # calculate the offset
            text_width, text_height = font.getbbox(watermark['text'])[2:4]  # Use getbbox() to get text size
            alpha = watermark.get('alpha', 128)  # Use alpha from the watermark dict, or default to 128 if not provided

            # calculate the position based on the specified position
            if watermark.get('position') == 'bottom':
                position = ((width - text_width) / 2, height - text_height - offset)
            elif watermark.get('position') == 'top':
                position = ((width - text_width) / 2, offset)
            elif watermark.get('position') == 'center':
                position = ((width - text_width) / 2, (height - text_height) / 2)
            else:
                position = ((width - text_width) / 2, height - text_height - offset)  # default to bottom

            # create a new image with an alpha channel (RGBA), draw the text on it, then composite that onto the original frame
            text_layer = Image.new('RGBA', frame.size, (255,255,255,0))  # fully transparent
            ImageDraw.Draw(text_layer).text(position, watermark['text'], (255, 255, 255, alpha), font=font)  # use alpha from the watermark dict
            frame = Image.alpha_composite(frame.convert('RGBA'), text_layer)
        frame_path = os.path.join(frame_dir, f"frame_{frame_number}.jpeg")
        frame = frame.convert("RGB")  # convert image to RGB mode
        frame.save(frame_path,format='jpeg')

    
        
    def open_image(self, image_path):
        try:
            with Image.open(image_path) as img:
                frame_number = int(os.path.split(image_path)[1].split('_')[1].split('.')[0])
                self.frame_queue.put((frame_number, img.copy()))
        except Exception as e:
            print(f"Error opening image {image_path}: {e}")


    def create_gif(self, frame_pattern, output_gif_path, frame_dir=None, frame_delays=None):
        if frame_dir is None:
            frame_dir = self.output_folder
        # image_paths = glob.glob(os.path.join(frame_dir, frame_pattern + '*.jpeg'))
        image_paths = glob.glob(os.path.join(frame_dir, f'{frame_pattern}*.jpeg'))
        image_paths.sort()

        # 清空队列，以防上次操作遗留数据
        while not self.frame_queue.empty():
            self.frame_queue.get()

        threads = []
        for image_path in image_paths:
            thread = threading.Thread(target=self.open_image, args=(image_path,))
            thread.start()
            threads.append(thread)

        for thread in threads:
            thread.join()

        if not image_paths:
            print(f"没有找到图片，请检查路径或模式。{os.path.join(frame_dir, frame_pattern + '*.png')}")
            return

        frames = [self.frame_queue.get()[1] for _ in range(len(image_paths))]
        frames[0].save(output_gif_path, save_all=True, append_images=frames[1:], duration=frame_delays, loop=0,optimize=True)



if __name__ == '__main__':
    # 使用示例
    watermark = {
        'text': "Generated by AI", 
        'font_scale': 1/30, 
        'position': 'bottom', 
        'offset_ratio': 0.02,
        'alpha': 190  # Set transparency to be 64 (0 is fully transparent, 255 is fully opaque)
    }
    gif_handler = GifHandler('./output_images')
    
    
    # frame_dir = gif_handler.download_and_split_gif("https://ty-huanlian.oss-cn-shanghai.aliyuncs.com/test/yato-herzog.gi")
    # print (frame_dir)
    # 要解析出这个图片,后面也是删除的!
    # 把gif合并成一张图
    gif_handler.create_gif('frame_', 'output.gif', frame_dir="/root/autodl-tmp/rope/utils/output_images/3ec781c55ece0cc01936a51fa5f9c083", duration=100)



